#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: yerunyuan@163.com
"""

import argparse
import os
import sys
import random

sys.path.append('.')

"""
   选出相同人脸的视频(限定每个人脸最多视频数)
"""


def read_merge_file(file_list):
    xiaoqu_faceid_subfaceids_list = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # gd/zs/dqjd/xgly:150100414a5444345203bcd04294b500_20190913_120347,150100414a5444345203bcd04294b500_20190923_141953,150100414a54443452067fa6d45c6600_20190927_075105,150100414a54443452067fa6d45c6600_20190913_163448:5
            lines = line.strip().split(':')
            if len(lines) > 2:
                serial_num = int(lines[2])
                if serial_num < 0:
                    continue
            xiaoqu = lines[0]
            faceids = lines[1].split(',')

            faceid_subfaceids = []
            for faceid in faceids:
                faceid_subfaceids.append(faceid.replace('_', ' ', 1))

            faceid_subfaceids_list = xiaoqu_faceid_subfaceids_list.get(xiaoqu, None)
            if faceid_subfaceids_list is None:
                xiaoqu_faceid_subfaceids_list[xiaoqu] = [faceid_subfaceids]
            else:
                faceid_subfaceids_list.append(faceid_subfaceids)
        print('Read Merge Count: {}'.format(file_count))
    return xiaoqu_faceid_subfaceids_list


def main():
    parser = argparse.ArgumentParser(description="Select Same Face Video 2")
    parser.add_argument(
        "--merge_file", default="/rootfs/media/kasim/DataSet/FaceCropImage/SelectMergeFaces.txt", help="path to merge file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/DataSet/FaceCropImage", help="path to output", type=str
    )
    parser.add_argument('--select_count', type=int, default=4, help='select count')
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    random.seed(123456)

    xiaoqu_faceid_subfaceids_list = read_merge_file(args.merge_file)

    out_noneed_file_path = os.path.join(output_dir, 'SelectFaceVideo.txt')
    out_noneed_file = open(out_noneed_file_path, 'w')

    for xiaoqu, faceid_subfaceids_list in xiaoqu_faceid_subfaceids_list.items():
        for faceid_subfaceids in faceid_subfaceids_list:
            file_names = []
            for faceid_subfaceid in faceid_subfaceids:
                device_id, video_name = faceid_subfaceid.split()
                file_name = os.path.join(xiaoqu, device_id, 'video', video_name + '.avi')
                file_names.append(file_name)
            if len(file_names) > args.select_count:
                random.shuffle(file_names)
                file_names = file_names[0:args.select_count]

            for file_name in file_names:
                out_noneed_file.write(file_name + '\n')

    out_noneed_file.close()
    os.system('chmod a+wr {}'.format(out_noneed_file_path))
    print('finish!')


if __name__ == '__main__':
    main()
